import numpy as np
import random
import math
import time
import tkinter as tk
from tkinter import messagebox
import matplotlib
matplotlib.rcParams['font.sans-serif'] = ['SimHei']
matplotlib.rcParams['axes.unicode_minus'] = False

class ReversiState:
    """反转棋游戏状态"""
    
    # 方向数组：8个方向
    DIRECTIONS = [(-1, -1), (-1, 0), (-1, 1), (0, -1), (0, 1), (1, -1), (1, 0), (1, 1)]
    
    def __init__(self, board_size=8):
        """初始化反转棋状态
        
        Args:
            board_size: 棋盘大小，默认为8x8
        """
        self.board_size = board_size
        self.board = np.zeros((board_size, board_size), dtype=int)
        self.current_player = 1  # 1为黑棋，2为白棋
        
        # 初始化棋盘，放置4个中间的棋子
        mid = board_size // 2
        self.board[mid-1, mid-1] = 2
        self.board[mid, mid] = 2
        self.board[mid-1, mid] = 1
        self.board[mid, mid-1] = 1
        
        # 游戏是否结束
        self.game_over = False
        
        # 缓存当前玩家的合法移动
        self._legal_moves = None
    
    def get_legal_moves(self):
        """获取当前玩家的合法移动"""
        if self._legal_moves is not None:
            return self._legal_moves
        
        moves = []
        for i in range(self.board_size):
            for j in range(self.board_size):
                if self.is_legal_move(i, j, self.current_player):
                    moves.append((i, j))
        
        self._legal_moves = moves
        return moves
    
    def is_legal_move(self, row, col, player):
        """判断玩家在指定位置落子是否合法"""
        # 如果位置已经有棋子，则不合法
        if self.board[row, col] != 0:
            return False
        
        # 检查8个方向是否能夹住对方的棋子
        opponent = 3 - player  # 对手的棋子编号（如果player=1，则opponent=2，反之亦然）
        
        for dr, dc in self.DIRECTIONS:
            r, c = row + dr, col + dc
            # 检查是否能找到对手的棋子
            if 0 <= r < self.board_size and 0 <= c < self.board_size and self.board[r, c] == opponent:
                # 继续沿着这个方向前进
                r += dr
                c += dc
                found_opponent = False
                
                while 0 <= r < self.board_size and 0 <= c < self.board_size:
                    if self.board[r, c] == 0:
                        break
                    if self.board[r, c] == player:
                        # 找到了自己的棋子，说明这个方向上可以夹住对手的棋子
                        return True
                    # 继续前进
                    r += dr
                    c += dc
        
        return False
    
    def make_move(self, row, col):
        """在指定位置落子，翻转对手的棋子"""
        if not self.is_legal_move(row, col, self.current_player):
            return False
        
        # 在指定位置放置棋子
        self.board[row, col] = self.current_player
        opponent = 3 - self.current_player
        
        # 翻转对手的棋子
        for dr, dc in self.DIRECTIONS:
            r, c = row + dr, col + dc
            to_flip = []
            
            # 寻找可能被翻转的棋子
            while 0 <= r < self.board_size and 0 <= c < self.board_size and self.board[r, c] == opponent:
                to_flip.append((r, c))
                r += dr
                c += dc
                
                # 如果找到了自己的棋子，说明这些对手的棋子可以被翻转
                if 0 <= r < self.board_size and 0 <= c < self.board_size and self.board[r, c] == self.current_player:
                    for fr, fc in to_flip:
                        self.board[fr, fc] = self.current_player
                    break
        
        # 切换玩家
        self.current_player = opponent
        
        # 清除缓存的合法移动
        self._legal_moves = None
        
        # 检查游戏是否结束
        self.check_game_over()
        
        return True
    
    def check_game_over(self):
        """检查游戏是否结束"""
        # 如果当前玩家没有合法移动，尝试切换玩家
        if not self.get_legal_moves():
            self.current_player = 3 - self.current_player
            self._legal_moves = None
            
            # 如果另一个玩家也没有合法移动，游戏结束
            if not self.get_legal_moves():
                self.game_over = True
                return True
        
        return False
    
    def get_winner(self):
        """获取游戏胜者"""
        if not self.game_over:
            return None
        
        # 计算黑白棋子数量
        black_count = np.sum(self.board == 1)
        white_count = np.sum(self.board == 2)
        
        if black_count > white_count:
            return 1
        elif white_count > black_count:
            return 2
        else:
            return 0  # 平局
    
    def get_score(self, player):
        """获取玩家的棋子数量"""
        return np.sum(self.board == player)
    
    def clone(self):
        """创建游戏状态的副本"""
        new_state = ReversiState(self.board_size)
        new_state.board = self.board.copy()
        new_state.current_player = self.current_player
        new_state.game_over = self.game_over
        return new_state

class MCTSNode:
    """蒙特卡罗树搜索节点"""
    
    def __init__(self, state, parent=None, move=None):
        """初始化MCTS节点
        
        Args:
            state: 游戏状态
            parent: 父节点
            move: 从父节点到达此节点的移动
        """
        self.state = state
        self.parent = parent
        self.move = move  # 从父节点到该节点的移动
        self.children = []  # 子节点列表
        self.wins = 0  # 胜利次数
        self.visits = 0  # 访问次数
        self.untried_moves = state.get_legal_moves()  # 未尝试的移动
        self.player = state.current_player  # 当前玩家
    
    def select_child(self, exploration_weight=1.0):
        """使用UCB1公式选择最有希望的子节点"""
        # 选择UCB值最高的子节点
        return max(self.children, key=lambda c: c.get_ucb(exploration_weight))
    
    def get_ucb(self, exploration_weight):
        """计算UCB1值"""
        if self.visits == 0:
            return float('inf')
        
        # UCB1公式 = win_rate + exploration_term
        exploitation = self.wins / self.visits
        exploration = exploration_weight * math.sqrt(2 * math.log(self.parent.visits) / self.visits)
        return exploitation + exploration
    
    def expand(self):
        """扩展一个新的子节点"""
        if not self.untried_moves:
            return None
        
        # 选择一个未尝试的移动
        move = random.choice(self.untried_moves)
        self.untried_moves.remove(move)
        
        # 创建新的游戏状态
        new_state = self.state.clone()
        new_state.make_move(move[0], move[1])
        
        # 创建子节点
        child = MCTSNode(new_state, parent=self, move=move)
        self.children.append(child)
        return child
    
    def update(self, result):
        """更新节点统计信息"""
        self.visits += 1
        # 如果获胜的玩家是该节点的玩家，增加胜利次数
        if result == self.player:
            self.wins += 1
    
    def is_fully_expanded(self):
        """检查是否所有可能的移动都已经尝试"""
        return len(self.untried_moves) == 0
    
    def is_terminal(self):
        """检查是否为终端节点（游戏结束）"""
        return self.state.game_over

class MCTS:
    """蒙特卡罗树搜索算法"""
    
    def __init__(self, time_limit=1.0, exploration_weight=1.0):
        """初始化MCTS算法
        
        Args:
            time_limit: 搜索时间限制(秒)
            exploration_weight: UCB公式中的探索权重
        """
        self.time_limit = time_limit
        self.exploration_weight = exploration_weight
    
    def search(self, state):
        """搜索最佳移动
        
        Args:
            state: 当前游戏状态
            
        Returns:
            最佳移动(row, col)
        """
        # 创建根节点
        root = MCTSNode(state)
        
        # 如果只有一个可能的移动，直接返回
        if len(root.untried_moves) == 1:
            return root.untried_moves[0]
        
        # 在时间限制内执行搜索
        end_time = time.time() + self.time_limit
        num_simulations = 0
        
        while time.time() < end_time:
            # 1. 选择
            node = self._select(root)
            
            # 2. 扩展
            if not node.is_terminal() and node.untried_moves:
                node = node.expand()
            
            # 3. 模拟
            if node:  # 可能在扩展阶段返回None
                result = self._simulate(node.state)
                
                # 4. 反向传播
                self._backpropagate(node, result)
            
            num_simulations += 1
        
        # 选择访问次数最多的子节点对应的移动
        best_child = max(root.children, key=lambda c: c.visits)
        
        print(f"MCTS搜索完成: 进行了{num_simulations}次模拟")
        
        # 打印每个移动的访问次数和胜率
        for child in root.children:
            move = child.move
            visits = child.visits
            win_rate = child.wins / visits if visits > 0 else 0
            print(f"移动 {move}: 访问次数 {visits}, 胜率 {win_rate:.2f}")
        
        return best_child.move
    
    def _select(self, node):
        """选择阶段: 从根节点选择最有希望的子节点"""
        while not node.is_terminal() and node.is_fully_expanded():
            node = node.select_child(self.exploration_weight)
        return node
    
    def _simulate(self, state):
        """模拟阶段: 从当前状态随机模拟到游戏结束"""
        # 创建状态的副本用于模拟
        state = state.clone()
        
        # 随机模拟直到游戏结束
        while not state.game_over:
            legal_moves = state.get_legal_moves()
            if not legal_moves:
                # 如果当前玩家没有合法移动，尝试切换玩家
                state.current_player = 3 - state.current_player
                state._legal_moves = None
                legal_moves = state.get_legal_moves()
                
                # 如果另一个玩家也没有合法移动，游戏结束
                if not legal_moves:
                    state.game_over = True
                    break
            
            # 随机选择一个合法移动
            move = random.choice(legal_moves)
            state.make_move(move[0], move[1])
        
        # 返回获胜的玩家
        return state.get_winner()
    
    def _backpropagate(self, node, result):
        """反向传播阶段: 更新节点统计信息"""
        while node:
            node.update(result)
            node = node.parent

class ReversiGame:
    """反转棋游戏界面"""
    
    def __init__(self):
        """初始化反转棋游戏"""
        self.state = ReversiState()
        self.ai = MCTS(time_limit=2.0)  # AI搜索时间限制为2秒
        self.ai_player = 2  # AI默认使用白棋
        
        # 设置GUI
        self.setup_gui()
    
    def setup_gui(self):
        """设置游戏GUI"""
        self.window = tk.Tk()
        self.window.title("反转棋 - MiniAlphaGo (MCTS)")
        self.window.geometry("700x600")
        
        # 顶部信息
        frame_info = tk.Frame(self.window)
        frame_info.pack(pady=10)
        
        self.label_turn = tk.Label(frame_info, text="当前玩家: 黑棋", font=("微软雅黑", 14))
        self.label_turn.pack(side=tk.LEFT, padx=20)
        
        self.label_score = tk.Label(frame_info, text="黑棋: 2 | 白棋: 2", font=("微软雅黑", 14))
        self.label_score.pack(side=tk.LEFT, padx=20)
        
        # 棋盘框架
        frame_board = tk.Frame(self.window)
        frame_board.pack(pady=10)
        
        self.canvas = tk.Canvas(frame_board, width=480, height=480, bg="green")
        self.canvas.pack()
        
        # 绘制棋盘
        self.draw_board()
        
        # 添加棋盘点击事件
        self.canvas.bind("<Button-1>", self.handle_click)
        
        # 控制面板
        frame_control = tk.Frame(self.window)
        frame_control.pack(pady=10)
        
        # AI难度选择
        frame_difficulty = tk.Frame(frame_control)
        frame_difficulty.pack(pady=5)
        
        tk.Label(frame_difficulty, text="AI思考时间:", font=("微软雅黑", 12)).pack(side=tk.LEFT)
        
        self.difficulty_var = tk.StringVar(value="2.0")
        difficulty_options = ["0.5", "1.0", "2.0", "3.0", "5.0"]
        difficulty_menu = tk.OptionMenu(frame_difficulty, self.difficulty_var, *difficulty_options, 
                                        command=self.change_difficulty)
        difficulty_menu.pack(side=tk.LEFT, padx=10)
        
        # 玩家选择
        frame_player = tk.Frame(frame_control)
        frame_player.pack(pady=5)
        
        tk.Label(frame_player, text="玩家棋色:", font=("微软雅黑", 12)).pack(side=tk.LEFT)
        
        self.player_var = tk.StringVar(value="黑棋")
        player_options = ["黑棋", "白棋"]
        player_menu = tk.OptionMenu(frame_player, self.player_var, *player_options, 
                                    command=self.change_player)
        player_menu.pack(side=tk.LEFT, padx=10)
        
        # 按钮
        frame_buttons = tk.Frame(frame_control)
        frame_buttons.pack(pady=10)
        
        self.btn_new_game = tk.Button(frame_buttons, text="新游戏", font=("微软雅黑", 12),
                                      command=self.new_game)
        self.btn_new_game.pack(side=tk.LEFT, padx=10)
        
        self.btn_ai_move = tk.Button(frame_buttons, text="AI走棋", font=("微软雅黑", 12),
                                     command=self.ai_make_move)
        self.btn_ai_move.pack(side=tk.LEFT, padx=10)
        
        self.btn_quit = tk.Button(frame_buttons, text="退出", font=("微软雅黑", 12),
                                 command=self.window.quit)
        self.btn_quit.pack(side=tk.LEFT, padx=10)
        
        # 状态栏
        self.status_text = tk.StringVar(value="游戏开始，黑棋先行")
        self.status_bar = tk.Label(self.window, textvariable=self.status_text, 
                                  font=("微软雅黑", 10), bd=1, relief=tk.SUNKEN, anchor=tk.W)
        self.status_bar.pack(side=tk.BOTTOM, fill=tk.X)
        
        # 如果AI是黑棋，则AI先走
        if self.ai_player == 1:
            self.window.after(500, self.ai_make_move)
    
    def draw_board(self):
        """绘制棋盘"""
        # 清除画布
        self.canvas.delete("all")
        
        # 绘制网格线
        cell_size = 480 // 8
        for i in range(9):
            # 水平线
            self.canvas.create_line(0, i * cell_size, 480, i * cell_size, fill="black", width=2)
            # 垂直线
            self.canvas.create_line(i * cell_size, 0, i * cell_size, 480, fill="black", width=2)
        
        # 绘制棋子
        for i in range(8):
            for j in range(8):
                if self.state.board[i, j] != 0:
                    self.draw_piece(i, j, self.state.board[i, j])
        
        # 绘制可能的移动位置
        if self.state.current_player != self.ai_player:
            legal_moves = self.state.get_legal_moves()
            for row, col in legal_moves:
                x = col * cell_size + cell_size // 2
                y = row * cell_size + cell_size // 2
                self.canvas.create_oval(x - 5, y - 5, x + 5, y + 5, fill="red")
        
        # 更新得分
        black_score = self.state.get_score(1)
        white_score = self.state.get_score(2)
        self.label_score.config(text=f"黑棋: {black_score} | 白棋: {white_score}")
        
        # 更新当前玩家
        player_name = "黑棋" if self.state.current_player == 1 else "白棋"
        self.label_turn.config(text=f"当前玩家: {player_name}")
    
    def draw_piece(self, row, col, player):
        """绘制棋子"""
        cell_size = 480 // 8
        x = col * cell_size + cell_size // 2
        y = row * cell_size + cell_size // 2
        color = "black" if player == 1 else "white"
        
        self.canvas.create_oval(x - cell_size // 2 + 5, y - cell_size // 2 + 5,
                               x + cell_size // 2 - 5, y + cell_size // 2 - 5,
                               fill=color, outline="black")
    
    def handle_click(self, event):
        """处理鼠标点击事件"""
        if self.state.game_over:
            return
        
        if self.state.current_player == self.ai_player:
            self.status_text.set("现在是AI的回合")
            return
        
        # 计算点击位置对应的棋盘坐标
        cell_size = 480 // 8
        col = event.x // cell_size
        row = event.y // cell_size
        
        if 0 <= row < 8 and 0 <= col < 8:
            if self.state.is_legal_move(row, col, self.state.current_player):
                self.state.make_move(row, col)
                self.draw_board()
                
                # 检查游戏是否结束
                if self.state.game_over:
                    self.game_over()
                    return
                
                # 如果下一个玩家是AI，让AI走棋
                if self.state.current_player == self.ai_player:
                    self.window.after(500, self.ai_make_move)
            else:
                self.status_text.set("无效的移动")
    
    def ai_make_move(self):
        """AI走棋"""
        if self.state.game_over or self.state.current_player != self.ai_player:
            return
        
        self.status_text.set("AI思考中...")
        self.window.update()
        
        # 让AI选择最佳移动
        start_time = time.time()
        move = self.ai.search(self.state.clone())
        elapsed_time = time.time() - start_time
        
        if move:
            self.state.make_move(move[0], move[1])
            self.draw_board()
            
            self.status_text.set(f"AI走棋: ({move[0]}, {move[1]}), 用时: {elapsed_time:.2f}秒")
            
            # 检查游戏是否结束
            if self.state.game_over:
                self.game_over()
        else:
            self.status_text.set("AI没有合法移动")
            
            # 如果AI没有合法移动，切换到玩家
            self.state.current_player = 3 - self.state.current_player
            self.state._legal_moves = None
            self.draw_board()
            
            # 检查游戏是否结束
            if not self.state.get_legal_moves():
                self.state.game_over = True
                self.game_over()
    
    def game_over(self):
        """游戏结束处理"""
        winner = self.state.get_winner()
        black_score = self.state.get_score(1)
        white_score = self.state.get_score(2)
        
        if winner == 0:
            message = f"游戏结束，平局！\n黑棋: {black_score}，白棋: {white_score}"
        elif winner == 1:
            message = f"游戏结束，黑棋胜利！\n黑棋: {black_score}，白棋: {white_score}"
        else:
            message = f"游戏结束，白棋胜利！\n黑棋: {black_score}，白棋: {white_score}"
        
        self.status_text.set(message)
        messagebox.showinfo("游戏结束", message)
    
    def new_game(self):
        """开始新游戏"""
        self.state = ReversiState()
        self.draw_board()
        
        self.status_text.set("新游戏开始，黑棋先行")
        
        # 如果AI是黑棋，则AI先走
        if self.ai_player == 1:
            self.window.after(500, self.ai_make_move)
    
    def change_difficulty(self, *args):
        """改变AI难度"""
        try:
            time_limit = float(self.difficulty_var.get())
            self.ai = MCTS(time_limit=time_limit)
            self.status_text.set(f"AI思考时间已设置为 {time_limit} 秒")
        except ValueError:
            self.status_text.set("无效的时间设置")
    
    def change_player(self, *args):
        """改变玩家棋色"""
        if self.player_var.get() == "黑棋":
            self.ai_player = 2  # AI使用白棋
            self.status_text.set("玩家使用黑棋，AI使用白棋")
        else:
            self.ai_player = 1  # AI使用黑棋
            self.status_text.set("玩家使用白棋，AI使用黑棋")
        
        # 开始新游戏
        self.new_game()
    
    def run(self):
        """运行游戏"""
        self.window.mainloop()

def main():
    """主函数"""
    game = ReversiGame()
    game.run()

if __name__ == "__main__":
    main()